import socket
import time
import threading

# 服务端 IP 和端口
SERVER_IP = '1.94.240.214'
SERVER_PORT = 9999

# 统计数据
disconnect_count = 0
received_message_count = 0
total_message_count = 0
latencies = []
heartbeat_message_count = 0

# 心跳检测函数，每隔一段时间向服务端发送心跳包
def send_heartbeat(sock):
    global heartbeat_message_count
    while True:
        time.sleep(45)  # 每 45 秒发送一次心跳消息
        try:
            sock.sendall(b"heartbeat")
            heartbeat_message_count += 1
        except:
            break

# 连接到服务端并接收消息
def connect_to_server():
    global disconnect_count, received_message_count, total_message_count, latencies

    while total_message_count < 50:  # 总消息数限制
        client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        client_socket.settimeout(180)  # 设置超时时间

        try:
            # 尝试连接到服务端
            client_socket.connect((SERVER_IP, SERVER_PORT))
            print(f"Connected to server {SERVER_IP}:{SERVER_PORT}")

            # 启动心跳线程
            heartbeat_thread = threading.Thread(target=send_heartbeat, args=(client_socket,))
            heartbeat_thread.daemon = True
            heartbeat_thread.start()

            while total_message_count < 2000:
                start_time = time.time()  # 记录发送时间

                try:
                    # 接收服务端的消息
                    message = client_socket.recv(1024)
                    if not message:
                        break

                    end_time = time.time()  # 记录接收时间
                    latency = end_time - start_time  # 计算消息时延
                    latencies.append(latency)

                    received_message_count += 1
                    total_message_count += 1
                    print(f"Received message {total_message_count}: {message.decode()}")
                except socket.timeout:
                    print("Connection timed out")
                    disconnect_count += 1  # 记录断链次数
                    break
                except Exception as e:
                    print(f"Error: {e}")
                    disconnect_count += 1  # 记录断链次数
                    break
        except Exception as e:
            print(f"Failed to connect: {e}")
        finally:
            client_socket.close()

        # 断开后等待 5 秒再尝试重连
        time.sleep(5)

# 主函数
if __name__ == "__main__":
    connect_to_server()

    # 打印统计结果
    print("\n--- 统计信息 ---")
    print(f"总消息数: {total_message_count}")
    print(f"成功接收到的消息数: {received_message_count}")
    print(f"断链次数: {disconnect_count}")
    print(f"发送的心跳消息数: {heartbeat_message_count}")

    if latencies:
        avg_latency = sum(latencies) / len(latencies)
        print(f"平均消息时延: {avg_latency:.6f} 秒")
        print(f"最小时延: {min(latencies):.6f} 秒")
        print(f"最大时延: {max(latencies):.6f} 秒")
    else:
        print("未收集到时延数据。")